from .base_strategy import BaseStrategy
import backtrader as bt
import backtrader.indicators as btind
import indicator.MyTTIndicator as Mytt
import indicator.myzhibiao as mmm
# 超级趋势指标
from indicator.superTrend import SuperTrend

###
# 130日线突破策略
# 使用4h级别, 唐齐安通道（130）, 
# 当最高价突破上轨后作为关键k，后续10根k线内，突破关键k的实体, 则开空单, 使用atr100波动率进行仓位管理
# 止损使用当前通道上轨止损，止盈，当价格突破下轨止盈
# 
# #
class MyStrategy(BaseStrategy):
    params = (
        ('donchian_period', 130),  # 唐齐安通道周期
        ('lookback', 10),  # 关键K线后观察周期
        ('atr_period', 14),      # ATR计算周期
        ('risk_exposure', 0.01),  # 风险暴露比例 (5%)
        # ('stop_loss_mult', 2.0),  # 止损倍数 (2倍ATR)
        # ('trail_mult', 3.0),      # 回撤止盈倍数 (3倍ATR)
        # ('take_profit_mult', 5.0), # 盈利止盈倍数 (5倍ATR)
    )

 
    def __init__(self):
        super().__init__()

        # 初始化指标
        self.donchian_high = btind.Highest(self.data.high(-1), period=self.p.donchian_period)
        self.donchian_low = btind.Lowest(self.data.low(-1), period=self.p.donchian_period)
        
        self.atr = btind.ATR(self.data, period=self.p.atr_period)
        
        self.supertrend = SuperTrend(
            period=10,
            multiplier=3
        )


         # 跟踪变量
        self.order = None
        self.key_bar_index = None
        self.breakout_type = None  # 'long' 或 'short'

        self.stop_loss_price = 0 # 止损价格
        self.take_profit_price = 0 # 止盈价格


        # 绘图相关的参数
        self.donchian_high.plotinfo.plot = True # 启用绘图，默认是启用的
        self.donchian_high.plotinfo.subplot = False  # 禁用子图（绘制在主图）
        self.donchian_low.plotinfo.subplot = False
        # 自定义悬停显示格式
        self.donchian_high.plotinfo.plotfmt = 'Donchian High: %.2f'  # 显示 2 位小数
        self.donchian_low.plotinfo.plotfmt = 'Donchian Low: %.2f'

    def start(self):
        pass
    def notify_order(self, order):
        super().notify_order(order)
        self.order = None

    def calculate_position_size(self):
        """计算固定风险比例的仓位大小
            本金 * 风险因子 * 现价 / atr
        """
        if self.atr[0] <=0:
            return 0
        account_value = self.broker.getvalue()
        # 风险金额
        risk_amount = account_value * self.p.risk_exposure
        return risk_amount / self.atr[0]

    def next(self):
        super().next()
        # print(f"时间：{self.data.datetime.datetime(0)}，收盘价：{self.data.close[0]}，最高价：{self.donchian_high[0]},最低：{self.donchian_low[0]}")
        # 有未完成的订单时, 不进行新订单
        if self.order:
            return
        # print(not self.position)
        if not self.position: #没有持仓
            if self.check_key_bar_conditions():
                return  # 已经记录关键K线，等待确认
            
            # 检查是否满足入场条件
            if self.check_entry_conditions():
                # 计算风险仓位
                size = self.calculate_position_size()
                if size == 0:
                    return
                
                if self.breakout_type == 'long':
                    self.stop_loss_price = self.donchian_low[0] # 使用最高价作为止损
                    self.take_profit_price = self.donchian_high[0] # 当前止盈价  实时变动
                    self.order = self.buy(size=size,price=self.data.close[0] * 1.01)
                elif self.breakout_type == 'short':
                    self.stop_loss_price = self.donchian_high[0] # 使用最低价作为止损
                    self.take_profit_price = self.donchian_low[0] # 最高价作为止盈
                    self.order = self.sell(size=size,price=self.data.close[0] * 0.99)
        else:# 有持仓
            # 自动计算止盈止损, 最高最低价等
            # self.update_stop_loss_take_profit()

            if self.check_exit_conditions():
                self.order = self.close()
                self.key_bar = None  # 重置关键K线
                self.breakout_type = None
                return


    def check_key_bar_conditions(self):
        '''检查是否出现关键K线'''
        # 突破下轨 - 潜在做多信号
        if self.data.low[0] < self.donchian_low[-1]:
            # 记录当前K线索引
            self.key_bar_index = len(self)
            self.breakout_type = 'long'
            # self.log(f'发现潜在做多关键K线,index:{self.key_bar_index} @{self.data.close[0]:.2f}')
            return True
        
        # 突破上轨 - 潜在做空信号
        elif self.data.high[0] > self.donchian_high[-1]:
            # 记录当前K线索引
            self.key_bar_index = len(self)
            self.breakout_type = 'short'
            # self.log(f'发现潜在做空关键K线,index:{self.key_bar_index} @{self.data.close[0]:.2f}')
            return True
        
        return False
    
    def check_entry_conditions(self):
        '''检查入场条件'''
        if self.key_bar_index is None:
            return False
        
        # 计算自关键K线以来的K线数量
        bars_since_key = len(self) - self.key_bar_index
        # 还在观察期内
        if bars_since_key <= self.p.lookback:
            temp_high = self.data.high.get(ago=-bars_since_key, size=1)[0]
            temp_low = self.data.low.get(ago=-bars_since_key, size=1)[0]
            # 做多条件：10根K线内收盘价低于关键K线最低价
            if (self.breakout_type == 'long' and 
                self.data.close[0] > temp_high):
                # self.log(f'做多信号确认 信号来源：前{bars_since_key}根k线，收盘价{self.data.close[0]:.2f}突破{temp_high}')
                return True
            
            # 做空条件：10根K线内收盘价高于关键K线最高价
            elif (self.breakout_type == 'short' and 
                self.data.close[0] < temp_low):
                # self.log(f'做空信号确认 信号来源：前{bars_since_key}根k线 收盘价{self.data.close[0]:.2f}跌破{temp_high}')
                return True
        
        # 超过观察期，重置关键K线
        elif bars_since_key > self.p.lookback:
            self.key_bar_index = None
            self.breakout_type = None
        
        return False
    
    def check_exit_conditions(self):
        '''检查出场条件（止损和止盈）'''
        if not self.position:
            return False
        
        self.take_profit_price = self.donchian_high[-1] if self.position.size >0 else self.donchian_low[-1]
        # 做多止损：价格跌破当前唐齐安通道下轨
        # if (self.position.size > 0 and 
        #     self.data.close[0] < self.stop_loss_price):
        #     self.log(f'做多止损 @{self.data.close[0]:.2f}')
        #     return True
        
        # 做多止盈：价格突破当前唐齐安通道上轨
        if (self.position.size > 0 and 
              self.data.high[0] > self.take_profit_price):
            self.log(f'做多止盈 @{self.data.close[0]:.2f}')
            return True
        
        # 做空止损：价格突破当前唐齐安通道上轨
        # if (self.position.size < 0 and 
        #       self.data.close[0] > self.stop_loss_price):
        #     self.log(f'做空止损 @{self.data.close[0]:.2f}')
        #     return True
        
        # 做空止盈：价格跌破当前唐齐安通道下轨
        if (self.position.size < 0 and 
              self.data.low[0] < self.take_profit_price):
            self.log(f'做空止盈 @{self.data.close[0]:.2f}')
            return True
        
        return False

    # 更新最高最低价格, 以及移动止盈价格
    def update_stop_loss_take_profit(self):
        if self.position:
            # 更新最高价和最低价
            if self.position.size > 0:  # 多头
                if self.highest_price == 0:
                    self.highest_price = self.data.close[0]
                else:
                    self.highest_price = max(self.highest_price, self.data.high[0])

                self.stop_loss_price = self.entry_price - self.p.stop_loss_mult * self.atr[0]
                
                # 回撤止盈
                trail_stop = self.highest_price - self.p.trail_mult * self.atr[0]
                self.stop_loss_price = max(self.stop_loss_price, trail_stop)
                
                # 盈利止盈
                self.take_profit_price = self.entry_price + self.p.take_profit_mult * self.atr[0]
            else:  # 空头
                if self.lowest_price == 0:
                    self.lowest_price = self.data.close[0]
                else:
                    self.lowest_price = min(self.lowest_price, self.data.low[0])

                self.stop_loss_price = self.entry_price + self.p.stop_loss_mult * self.atr[0]
                
                # 回撤止盈
                trail_stop = self.lowest_price + self.p.trail_mult * self.atr[0]
                self.stop_loss_price = min(self.stop_loss_price, trail_stop)
                
                # 盈利止盈
                self.take_profit_price = self.entry_price - self.p.take_profit_mult * self.atr[0]